Radio - Rate Distortion Optimization for LLM Compression
Table of Contents
Radio provides a novel view for quantization with optimization technique and rate-distortion theory, as well as a closed form for quantization bit. The drawback is that it does not consider hardware efficiency.
We model the next-token prediction as a single function: \(\mathbf{Z}=f(\mathbf{X}; \Theta_0, \Theta_1, \dots, \Theta_n, \Theta_{n+1})\). The model has \(n\) layers where \(\Theta_0\) is word embedding and \(\Theta_{n+1}\) is the prediction head.
1. RADIO Solution
We use the settings including next token prediction. We use follow quantization schema, suppose bit depth is \(B\) and step size is \(D\), the dequantization can be expressed as
\[ \theta^q(B,D) = D \cdot \Big( \texttt{clip}\big( \texttt{ceil}(\theta/D), -2^{B-1}, 2^{B-1}-1 \big) + 0.5 \Big) \]
In RADIO, we’ll use one pair of \((B,D)\) to quantize a small group of weights.
1.1. TODO Bit Depth Assignment
Suppose we want to quantize matrix \(\Theta_n\) with \(P_n\) parameters. Following the intuition that parameters that are more sensitive to output distortion should be allocated more bits, we can formalize the optimization problem:
\begin{array}{rlllllll} \min & d(\{ B_n \}) = \mathbb{E}_{\mathbf{X}} \Vert f(\mathbf{X}; \{ \Theta^q_n(B_n) \}) - f(\mathbf{X}) \Vert^2 \\ \text{s.t.} & r(\set{B_n}) = \sum_{i=1}^N P_nB_n - R(\sum_{i=1}^N P_n) = 0 \end{array}- We use \(D_n^\ast(B_n)\) to denote the step size given \(B_n\). \(\Theta^q_n (B_n)\) is the abbreviation for \(\Theta^q_n(B_n,D_n^\ast(B_n))\), which means the quantized weight matrix \(\Theta_n^q\) w.r.t. \(\Theta_n\).
- \(R\) is the target compression bit width, which is predefined hyperparameter.
- The object function comes from rate-distortion theory. In short, we try to minimize the expected gap between original model and quantized model. Intuitively, if the quantized model can mimic the behaviour well, then we may say the preformance degradation is not large.
Then, we want to solve for the design variables \(B_i\). We can apply the Lagrangian, and we get
\begin{equation} \frac{1}{P_n} \frac{\partial d(\set{B_n})}{\partial B_n} = -\lambda, \end{equation} \begin{equation} r(\set{B_n})=0 \end{equation}The next step, is to find an approximation of the partial derivative. An classic result in rate distortion theory1 gives that for any random variable with finite variance, its quantization error halves with every additional bit at a sufficiently high bit depth. In this problem setting, we have
\begin{equation} \begin{split} -\frac{1}{2\ln 2} \frac{\partial d(\set{B_n})}{\partial B_n} & \approx \mathbb{E}_{\mathbf{X}} \Big\Vert \frac{\partial f(\set{\Theta^q_n(B_n)})}{\partial \Theta_n} \Delta^q (B_n) \Big\Vert ^2 \\ & \approx P_n H_n G_n^2 S_n^2 2^{-2B_n} \\ & \overset{\text{def}}{=} d_n(B_n) \end{split} \end{equation}- \(G_n^2\) represents the variances of elements of \(\partial_{\Theta_n} f(\mathbf{X})\)
- \(S_n^2\) represents the variances of elements \(\Theta^q_n\)
- \(H_n\) is a quantization coefficient that depends only on the type of weight distribution1.
After approximating the gradient, we can apply dual ascent2 to alternatively update \(B_n\) and \(\lambda\).
\begin{equation} \begin{split} B_n &\gets \texttt{clamp}\Big( \frac{1}{2}\log_2\big( \frac{G_n^2 S_n^2}{V/2 \ln 2} \big), 0, B_{\max}=8 \Big) \\ V &\gets V+\beta \Big( r(\set{B_n}) \Big) \end{split} \end{equation}- \(\beta\) represents a step size for dual update.
1.2. TODO Quantization Step Size
Footnotes:
Allen Gersho, and Robert M. Gray. Vector Quantization and Signal Compression. Kluwer, Norwell, MA, USA 1991.
Stephen Boyd, Neal Parikh, Eric Chu, Borja Peleato, and Jonathan Eckstein. Distributed optimization and statistical learning via the alternating direction method of multipliers. Found. Trends® Mach. Learn., 3(1):1122, 2011.